import os
import time
import json
import subprocess
import pandas as pd
from ultralytics import YOLO

RESULTS = []

def train_yolov5():
    print("\n=== YOLOv5 Training ===")
    start = time.time()
    command = [
        "python", "train.py",
        "--img", "640",
        "--batch", "16",
        "--epochs", "100",
        "--data", "../data.yaml",  # 假设 run_training.py 在 yolov5 目录外
        "--weights", "yolov5n.pt",
        "--device", "mps"
    ]
    subprocess.run(command, cwd="./yolov5")
    end = time.time()
    
    result_path = "./yolov5/runs/train/exp/results.json"
    if os.path.exists(result_path):
        with open(result_path) as f:
            metrics = json.load(f)
        RESULTS.append({
            "Model": "YOLOv5n",
            "mAP50": metrics.get("metrics/mAP_0.5", None),
            "mAP50-95": metrics.get("metrics/mAP_0.5:0.95", None),
            "Time(s)": round(end - start, 2)
        })
    else:
        RESULTS.append({"Model": "YOLOv5n", "mAP50": "N/A", "mAP50-95": "N/A", "Time(s)": round(end - start, 2)})

def train_yolov8():
    print("\n=== YOLOv8 Training ===")
    start = time.time()
    model = YOLO("yolov8n.pt")
    results = model.train(data="data.yaml", epochs=100, imgsz=640, device="mps")
    end = time.time()
    
    metrics = results.metrics
    RESULTS.append({
        "Model": "YOLOv8n",
        "mAP50": round(metrics.get("metrics/mAP50", 0), 4),
        "mAP50-95": round(metrics.get("metrics/mAP50-95", 0), 4),
        "Time(s)": round(end - start, 2)
    })

def train_yolov11():
    print("\n=== YOLOv11 Training (YOLOv8n from v11.0+) ===")
    start = time.time()
    model = YOLO("yolov8n.pt")  # 当前版本仍用 yolov8n.pt，ultralytics v11.0+
    results = model.train(data="data.yaml", epochs=100, imgsz=640, device="mps")
    end = time.time()
    
    metrics = results.metrics
    RESULTS.append({
        "Model": "YOLOv11n",
        "mAP50": round(metrics.get("metrics/mAP50", 0), 4),
        "mAP50-95": round(metrics.get("metrics/mAP50-95", 0), 4),
        "Time(s)": round(end - start, 2)
    })

if __name__ == "__main__":
    train_yolov5()
    train_yolov8()
    train_yolov11()

    df = pd.DataFrame(RESULTS)
    print("\n=== Final Results ===")
    print(df.to_string(index=False))

    df.to_csv("training_comparison.csv", index=False)
    print("\nSaved to training_comparison.csv ✅")
